import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.pyplot import scatter
from pandas.core.internals.construction import rec_array_to_mgr
from sklearn.cluster import KMeans
from sklearn.impute import KNNImputer
from sklearn.metrics import silhouette_score

data = pd.read_csv("customers.csv")
x = data.iloc[:,[3,4]]

sse_list =  []
sc_list = []
for k in range(2,20):
    model = KMeans(n_clusters=k)
    model.fit(x)
    sse_list.append(model.inertia_)

    y_pre = model.predict(x)
    sc_list.append(silhouette_score(x,y_pre))

# t1 = plt.subplot(1,2,1)
# t1.plot(range(2,20),sse_list,marker="o")
# t2 = plt.subplot(1,2,2)
# t2.plot(range(2,20),sc_list,marker="o")
# plt.show()

model = KMeans(n_clusters=5)
model.fit(x)
y_pre = model.predict(x)

plt.scatter(x.values[y_pre == 0, 0], x.values[y_pre == 0, 1], s=100, c='green', label='Normal')
plt.scatter(x.values[y_pre == 1, 0], x.values[y_pre == 1, 1], s=100, c='red', label='Standard')
plt.scatter(x.values[y_pre == 2, 0], x.values[y_pre == 2, 1], s=100, c='blue', label='Traditional')
plt.scatter(x.values[y_pre == 3, 0], x.values[y_pre == 3, 1], s=100, c='black', label='Youth')
plt.scatter(x.values[y_pre == 4, 0], x.values[y_pre == 4, 1], s=100, c='yellow', label='TA')
plt.scatter(model.cluster_centers_[:, 0], model.cluster_centers_[:, 1], s=300, c='black', label='Centroids')

plt.title('Clusters of customers')
plt.xlabel('Annual Income (k$)')
plt.ylabel('Spending Score (1-100)')
plt.legend()
plt.show()





